
import re,pickle
# 用标注数据作为验证集看模型效果,只关注mod类型的语义角色

def process(path):
    all_role = set()
    event_role = ['eCoo','eSelt','eEqu','ePrec','eSucc','eProg','eCau','eAdvt','eResu','eInf','eCond','eSupp','eConc','eAban','eMetd','ePurp','ePref','eSum','eRect']
    with open(path, 'r', encoding='utf-8') as file:
        content = file.readlines()
        print(len(content),content[2])
        result_sent = []
        for i in range(0,len(content)-1,3):
            sent, record, idx = [], [], []
            head = content[i].strip().split()
            head = [key.split(']')[1] for key in head]
            pos = head[1:]
            fenci = []
            for key in pos:
                s= key.split('/')
                fenci.append(key[:len(key)-len(s[-1])-1])
            # fenci = [key.split('/')[0] for key in pos]
            # print('fenci:', fenci)
            num,s=[],0   
            for k in fenci:
                num.append(list(range(s,s+len(k))))
                s=s+len(k)
            sent.append(''.join(i for i in fenci))
            sent.append(num)
            role = content[i+1].split()
            # print(role)
            s = 0
            for item in role:  # 找到核心谓词节点id:s
                id = re.split(r'[_\(]', item)
                # print(id)  # ['[4]企业', '[3]投资', 'rPat)']
                h_id, h_w = int(id[0].strip('[').split(']')[0])-1, id[0].strip('[').split(']')[1]
                t_id, t_w = int(id[1].strip('[').split(']')[0])-1, id[1].strip('[').split(']')[1]
                rel = id[2].strip(')')
                if rel == "Root":
                    s = t_id
                    break
            other_v = []
            for item in role:  # 找到跟核心谓词节点形成事件关系的节点id]
                id = re.split(r'[_\(]', item)
                # print(id)  # ['[4]企业', '[3]投资', 'rPat)']
                h_id, h_w = int(id[0].strip('[').split(']')[0])-1, id[0].strip('[').split(']')[1]
                t_id, t_w = int(id[1].strip('[').split(']')[0])-1, id[1].strip('[').split(']')[1]
                rel = id[2].strip(')')
                if rel in event_role:
                    other_v.append(h_id)
                    other_v.append(t_id) # 存储非核心谓词的id
                    break
            other_v.append(s)
            for item in role:
                id = re.split(r'[_\(]', item)
                h_id, h_w = int(id[0].strip('[').split(']')[0])-1, id[0].strip('[').split(']')[1]
                t_id, t_w = int(id[1].strip('[').split(']')[0])-1, id[1].strip('[').split(']')[1]
                rel = id[2].strip(')')
                if h_id<0 or t_id<0 or rel == 'mPunc':
                    continue
                if rel in sbj:
                    rel = sbj.get(rel,rel)
                if rel in obj:
                    rel = obj.get(rel,rel)
                if h_id in other_v and rel not in event_role and rel[0] not in ['r','d','m'] and rel not in r5:
                    all_role.add(rel)
                    idx.append((h_id, rel, t_id))
                    record.append((h_w,rel,t_w))
                # elif t_id in other_v:
                #     idx.append((t_id, rel, h_id))
                #     record.append((t_w,rel,h_w))
            sent.append(idx)
            sent.append(record)
            sent.append(pos)
            result_sent.append(sent)
    # print(len(result_sent))  # 10068
    # print(result_sent[:2])
    # print(len(all_role))
    # print(all_role)
    return result_sent

if __name__=="__main__":
    sbj = {'Agt':'sbj_Agt', 'Exp':'sbj_Exp', 'Aft':'sbj_Aft', 'Poss':'sbj_Poss'}
    obj = {'Pat':'obj_Pat', 'Cont':'obj_Cont', 'Prod':'obj_Prod', 'Orig':'obj_Orig', 'Comt':'obj_Comt', 'Comp':'obj_Comp'}
    r5 = {'Quan':'Quantity','Qp':'Quantity phrase','Freq':'Frequency','Seq':'Sequence','Nvar':'Variable','Nini':'Initial Number','Nfin':'Final Number'}

    train_sent = process('CTB_data/News/news.train.flat.txt')
    valid_sent = process('CTB_data/News/news.valid.flat.txt')
    test_sent = process('CTB_data/News/news.test.flat.txt')
    print('train_sent:', len(train_sent))
    print('valid_sent:', len(valid_sent))
    print('test_sent:', len(test_sent))
    
    train = 'role/'+ 'train_data.pkl'
    val = 'role/'+ 'val_data.pkl'
    test = 'role/'+ 'test_data.pkl'
    with open(train,'wb') as out_file:
        pickle.dump(train_sent,out_file) 
    with open(val,'wb') as out_file:
        pickle.dump(valid_sent,out_file) 
    with open(test,'wb') as out_file:
        pickle.dump(test_sent,out_file) 



'''
make_data
batch[3] :'抓住这一环，就为提高产品质量打下了基础。'
batch[4] :[[0, 1], [2], [3], [4], [5], [6], [7], [8, 9]] 
batch[5] :[(8, 'sbj_HX', 6), (8, 'obj_HX', 14),(8, '_BY', 14)]
batch[6] :[('抓住', 'obj', '环'), ('打下', 'obj', '基础')]
batch[7] :['抓住/v', '这/r', '一/m', '环/n', '，/w', '就/d', '为/p', '提高/v', '产品/n', '质量/n', '打下/v', '了/y', '基础/n', '。/w']
'''